#!/usr/bin/env python3
# stitch_mt.py
import os
import glob
import argparse
import concurrent.futures as futures
from typing import List, Tuple

import cv2
import numpy as np



"""
python3 pinjie.py \
  --img-dir "/home/chengjinlai/qiqi-guaiguai-xiaowanyi/opencv-图像拼接/jpg" \
  --pattern "*.jpg" \
  --threads 20 \
  --max-side 1800 \
  --mode panorama \
  --out result.jpg

"""
cv2.setNumThreads(os.cpu_count())
print("OpenCV threads:", cv2.getNumThreads())


def load_and_preprocess(path: str, max_side: int = 0) -> Tuple[str, np.ndarray]:
    """
    读取并可选等比例缩放一张图像。
    返回 (文件路径, 图像)；若读取失败则图像为 None。
    """
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    if img is None:
        return path, None

    if max_side and max(img.shape[:2]) > max_side:
        h, w = img.shape[:2]
        scale = max_side / float(max(h, w))
        nh, nw = int(round(h * scale)), int(round(w * scale))
        img = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_AREA)
    return path, img


def read_images_multithreaded(img_files: List[str], workers: int, max_side: int) -> List[np.ndarray]:
    images = []
    failed = []
    with futures.ThreadPoolExecutor(max_workers=workers) as ex:
        for path, img in ex.map(lambda p: load_and_preprocess(p, max_side), img_files):
            if img is None:
                failed.append(path)
            else:
                images.append(img)

    if failed:
        print("以下图片读取失败，将被跳过：")
        for p in failed:
            print("  -", p)
    return images


def build_stitcher(mode: str = "panorama") -> cv2.Stitcher:
    """
    mode: 'panorama' 或 'scans'
    """
    mode_const = cv2.Stitcher_PANORAMA if mode.lower() == "panorama" else cv2.Stitcher_SCANS
    try:
        stitcher = cv2.Stitcher_create(mode_const)  # OpenCV 4
    except AttributeError:
        stitcher = cv2.createStitcher(mode_const)   # 兼容旧版本
    return stitcher


def main():
    parser = argparse.ArgumentParser(description="多线程读取 + OpenCV 并行拼接")
    parser.add_argument("--img-dir", default="/home/chengjinlai/qiqi-guaiguai-xiaowanyi/opencv-图像拼接/jpg",
                        help="待拼接图片目录")
    parser.add_argument("--pattern", default="*.jpg", help="文件通配符，比如 *.jpg / *.png")
    parser.add_argument("--threads", type=int, default=os.cpu_count() or 4,
                        help="Python 读取线程数，同时也设置 OpenCV 线程数")
    parser.add_argument("--max-side", type=int, default=0,
                        help="预处理时最长边等比例缩放到该像素（0 表示不缩放）")
    parser.add_argument("--mode", choices=["panorama", "scans"], default="panorama",
                        help="Stitcher 模式：panorama(全景) / scans(平面扫描)")
    parser.add_argument("--out", default="result.jpg", help="输出文件名（保存在图片目录下）")
    args = parser.parse_args()

    img_dir = args.img_dir
    img_files = sorted(glob.glob(os.path.join(img_dir, args.pattern)))
    if not img_files:
        print("未找到任何图片，请检查路径与通配符：", os.path.join(img_dir, args.pattern))
        return

    print(f"共找到 {len(img_files)} 张图片，开始多线程读取（workers={args.threads}，max_side={args.max_side}）…")

    # 配置 OpenCV 自身的线程数（影响特征提取/匹配/变换等原生计算）
    try:
        cv2.setNumThreads(max(1, int(args.threads)))
    except Exception:
        pass  # 某些构建可能不支持

    # 可选：开启/关闭 OpenCL（有些环境开启会更快，有些反而慢，按需切换）
    try:
        cv2.ocl.setUseOpenCL(True)
    except Exception:
        pass

    # 1) 多线程读取 + 预处理
    images = read_images_multithreaded(img_files, workers=args.threads, max_side=args.max_side)
    print(f"成功读取 {len(images)} 张图片。")

    if len(images) < 2:
        print("需要至少 2 张图片才能拼接。")
        return

    # 2) 创建 Stitcher（内部已有原生并行）
    stitcher = build_stitcher(args.mode)

    # 如需更细的控制，可设置一些参数（不同 OpenCV 版本可用项略有差异）
    try:
        # 示例：提升鲁棒性（代价是速度）
        stitcher.setPanoConfidenceThresh(0.6)
        stitcher.setWaveCorrection(True)
    except Exception:
        pass

    print("开始拼接（由 OpenCV 原生并行负责加速）…")
    status, pano = stitcher.stitch(images)

    # 3) 结果处理
    if status != cv2.Stitcher_OK:
        print("拼接失败！错误码:", status)
        return

    out_path = os.path.join(img_dir, args.out)
    ok = cv2.imwrite(out_path, pano)
    if not ok:
        print("写入结果失败：", out_path)
        return

    print("拼接完成，结果保存在:", out_path)


if __name__ == "__main__":
    main()
